/************************************************************/
/*                Utility functions                         */
/*              Author:  Wayne Zhang                        */
/*            actuary_zhang@hotmail.com                     */
/************************************************************/

/**
 * @file utilities.c
 * @brief Utility functions for simple arithmetic calculation,
 * matrix computation and numerical derivatives. 
 * @author Wayne Zhang                         
 */

/* headers for R related definitions */
#include <R.h>
#include <Rinternals.h>
#include <Rmath.h>
#include <R_ext/Lapack.h>        /* for BLAS and Lapack related */
#include "utilities.h"


/** Allow for translation of error messages */
#ifdef ENABLE_NLS		
#include <libintl.h>
#define _(String) dgettext ("cplm", String)
#else
#define _(String) (String)
#endif

/** zero an array */
#define AZERO(x, n) {int _I_, _SZ_ = (n); for(_I_ = 0; _I_ < _SZ_; _I_++) (x)[_I_] = 0;}

/** constant used in numerical differentiation */
#define DIFF_EPS 0.001


/************************************************************/
/*               Simple arithmetic utility                  */
/************************************************************/

/*
 * Compute sample medians 
 *
 * @param arr vector of sample values 
 * @param n length of the vector
 */

double median(double *x, int n) {
  double *xs = Calloc(n, double); 
  Memcpy(xs, x, n); 
  R_rsort(xs, n);
  if (n % 2 == 0)
    return((xs[n/2] + xs[n/2 - 1])/2);
  else 
    return(xs[n/2]); 
}


/**
 * Compute sample variance 
 *
 * @param x samples in long vector 
 * @param n number of samples
 *
 * @return sample variance
 */

double var(double *x, int n){
  double ans = 0.0, m = mean(x, n) ;
  for (int i = 0; i < n; i++)
    ans += (x[i] - m) * (x[i] - m);
  ans /= n - 1.0  ;
  return ans;
}

/**
 * Compute sample covariance matrix 
 *
 * @param n number of samples
 * @param p number of variables (columns), p >2
 * @param x samples in long vector 
 * @param ans vector to store computed covariance matrix
 *
 */
void cov(int n, int p, double *x, double *ans){
  double *one = Calloc(n * n, double),
    *x2 = Calloc(n * p, double),
    *x3 = Calloc(n * p, double);
  double alpha = -1.0 / n, beta = 1.0, beta2 = 0.0;

  // subtract mean    
  for (int i = 0; i < n * n; i++) one[i] = 1.0 ;
  Memcpy(x2, x, n * p) ;
  Memcpy(x3, x, n * p); 
  F77_CALL(dgemm)("N", "N", &n, &p, &n, &alpha, one,
		  &n, x2, &n, &beta, x3, &n);
  Memcpy(x2, x3, n * p) ;
  AZERO(ans, p * p) ;
    
  // compute covariance 
  F77_CALL(dgemm)("T", "N", &p, &p, &n, &beta, x2,
		  &n, x3, &n, &beta2, ans, &p);
  for (int i = 0; i < p * p; i++) ans[i] /= (n - 1) * 1.0 ;
  Free(one) ;
  Free(x2) ;
  Free(x3) ;
}




/************************************************************/
/*                 Matrix computations                      */
/************************************************************/


/** 
 * column or row sum of a matrix 
 *
 * @param m number of rows
 * @param n number of columns
 * @param margin indicator of row sum (margin = 1) or column sum (margin = 2)
 * @param x the matrix to be summed 
 * @param ans output vector 
 * 
*/
void sum_mat(int m, int n, int margin, double *x, double *ans){
  // row sums
  if (margin == 1){
    for (int i = 0; i < m; i++){
      ans[i] = 0;
      for (int j = 0; j < n; j++){
	ans[i] += x[i + j * m];
      }
    }
  } else  //col sums
    if (margin == 2){
      for (int i = 0; i < n; i++){
	ans[i] = dcumsum(x, m);
	x += m;
      }
    } else {
      error("Wrong dimensions in sum_mat.\n");   
    }
}


/**
 * Multiply a matrix and a vector 
 *
 * @param trans transpose of matrix?
 * @param m row count of matrix
 * @param n column count of matrix
 * @param A input matrix
 * @param x input vector
 * @param out output vector 
 *
 */
void mult_mv(char *trans, int m, int n, double *A,
             double *x, double *out){
  double one = 1.0, zero = 0.0 ;
  int incx = 1;
  F77_CALL(dgemv)(trans, &m, &n, &one, A, &m, x, &incx,
		  &zero, out, &incx) ;
}


/**
 * compute t(x) * x
 *
 * @param m row dimension of the matrix
 * @param n column dimension of the matrix
 * @param x the input matrix  
 * @param out output results
 *
 */

void mult_xtx(int m, int n, double *x, double *out){
  double alpha = 1.0, beta = 0.0, *x2 = Calloc(m * n, double);
  Memcpy(x2, x, m * n) ;
  F77_CALL(dgemm)("T", "N", &n, &n, &m, &alpha, x2, &m,
		  x, &m, &beta, out, &n) ;
  Free(x2) ;
}

/**
 * compute the lower cholesky factor
 *
 * @param d dimension of the matrix
 * @param v input matrix
 * @param iv output cholesky factor
 *
 */
void chol(int d, double *v, double *iv){    
  int info = 0;
  // cholesky factor of v
  Memcpy(iv, v, d * d) ;   
  F77_CALL(dpotrf)("L", &d, iv, &d, &info) ;
  if (info)  error(_("Error %d in Cholesky decomposition."), info) ;   
}

/**
 * invert a positive symmetric matrix 
 *
 * @param d dimension of the matrix
 * @param v input matrix
 * @param iv output inverse of the matrix
 *
 */
int solve_po(int d, double *v, double *iv){    
  int info = 0;
  // cholesky factor of v
  chol(d, v, iv) ;
  // compute inverse    
  F77_CALL(dpotri)("L", &d, iv, &d, &info) ;    
  if (info) {
    return(info);
  } else {
    // fill upper triangle 
    for (int i = 0; i < d - 1; i++){
      for (int j = i + 1; j < d; j++)
	iv[j * d + i] = iv[i * d + j] ;
    }
    return(0);
  }   
}


/**
 * Solve the equation Ax=b where A is positive definite
 *
 * @param d dimension of the matrix
 * @param A input pos def matrix
 * @param b input vector
 * @param x output vector
 *
 */
int solve_po2(int d, double *A, double *b, double *x){
  double *iA = Calloc(d * d, double);
  int error = solve_po(d, A, iA);
  if (error){
    Free(iA);
    return(error);
  } else{
    int d2 = d;
    mult_mv("N", d, d2, iA, b, x);
    Free(iA);
    return(0);
  }
}

/************************************************************/
/*                Compute numerical derivatives             */
/************************************************************/


/**
 * Compute numerical gradient 
 *
 * @param n length of parmaters
 * @param x values at which to evaluate the gradient
 * @param myfunc user specified function 
 * @param data struct used in myfunc
 * @param ans vector to store the gradient 
 *
 */

void grad(int n, double *x,  double (*myfunc)(double *x, void *data), 
          void *data, double *ans){
  double y1, y2 ;
  for (int i = 0; i < n; i++){
    x[i] += DIFF_EPS ;
    y1 = myfunc(x, data) ;
    x[i] -= 2 * DIFF_EPS ;        
    y2 = myfunc(x, data) ;
    ans[i] = (y1 - y2) / DIFF_EPS * 0.5 ;
    x[i] += DIFF_EPS ;
  }
}


/**
 * Compute numerical hessian matrix  
 *
 * @param n length of parmaters
 * @param x values at which to evaluate the hessian
 * @param myfunc user specified function 
 * @param data struct used in myfunc
 * @param ans n*n vector to store the hessian matrix 
 *
 */
void hess(int n, double *x, double (*myfunc)(double *x, void *data), 
          void *data, double *ans){
  double *y1 = Calloc(n, double),
    *y2 = Calloc(n, double)  ;
  for (int i = 0; i < n; i++){
    x[i] += DIFF_EPS ;
    grad(n, x, myfunc, data, y1) ;
    x[i] -= 2 * DIFF_EPS ;
    grad(n, x, myfunc, data, y2) ;
    for (int j = 0; j < n; j++)
      ans[j + i * n] = (y1[j] - y2[j]) / DIFF_EPS * 0.5 ;
    x[i] += DIFF_EPS ;
  }
  Free(y1) ; Free(y2) ;
}

